import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms

class FeatureExtractor(nn.Module):
    def __init__(self, model, target_layers):
        super(FeatureExtractor, self).__init__()
        self.model = model
        #print(self.model)
        self.target_layers = target_layers
        self.features = {}
        self._register_hooks()

    def _register_hooks(self):
        def hook_fn(module, input, output, key):
            self.features[key] = output

        for name, module in self.model.named_modules():
            #print(name)
            if name in self.target_layers:
                module.register_forward_hook(lambda m, i, o, k=name: hook_fn(m, i, o, k))
            

    def forward(self, x):
        self.features = {}
        self.model(x)
        return self.features
    


class ResizeFeatureExtractor(nn.Module):
    def __init__(self, model, size, target_layers):
        super(ResizeFeatureExtractor, self).__init__()
        self.resize = transforms.Resize(size)
        self.model = model
        self.target_layers = target_layers
        self.features = {}
        self._register_hooks()

    def _register_hooks(self):
        def hook_fn(module, input, output, key):
            self.features[key] = output

        for name, module in self.model.named_children():
            if name in self.target_layers:
                module.register_forward_hook(lambda m, i, o, k=name: hook_fn(m, i, o, k))

    def forward(self, x):
        self.features = {}
        x = self.resize(x)
        self.model(x)
        return self.features